//
// Created by y z l on 2023/7/22.
//
#include <iostream>
#include <vector>
#include <algorithm>
#include <string>
using namespace std;

class Solution {
public:
    vector<int> findDiagonalOrder(vector< vector<int> >& mat) {
        int iter = 0;
        int len = mat.size() * mat[0].size();
        int col_len = mat[0].size();
        int raw_len = mat.size();
        vector<int> ans_vec(len,0);
        int i=0,j=0;
        while(iter < len){
            // 左上
            while(iter < len && j < col_len && i>=0){
                ans_vec[iter] = mat[i][j];
                iter++;
                i--;
                j++;
            }
            if(j >= col_len){
                j = col_len - 1;
                i+=2;
            } else if(i<0){
                i = 0;
            }
            // 右下
            while(iter < len && i < raw_len && j >= 0){
                ans_vec[iter] = mat[i][j];
                iter++;
                i++;
                j--;
            }
            if(i >= raw_len){
                i = raw_len - 1;
                // 2 1 -> 3 0
                j+=2;
            } else if(j < 0){
                j = 0;
            }
        }
        return ans_vec;
    }
};

int main(){
    vector< vector<int> > matrix;
    int num=1;
    for(int i=0;i<3;i++){
        vector<int> vec;
        for(int i=0;i<3;i++){
            vec.push_back(num++);
        }
        matrix.push_back(vec);
    }
    // matrix[1][1] = 0;
    Solution solu;
    vector<int> ans_vec = solu.findDiagonalOrder(matrix);
    // solu.setZeroes(matrix);

    return 0;
}